# coding=utf-8
import numpy

'''
朴素贝叶斯算法实现
'''


class Bayes:

    def __init__(self):

        # 数据集涉及的特征数量，初始化为1
        self.feature_length = 1

        # 各类别的概率['类别1':概率1, '类别2':概率2, ......]
        self.label_probability_dict = dict()

        # 以字典的方式存储类别与特征向量
        # 格式为{'类别1':[特征向量1, 特征向量2, ......], '类别2':[特征向量1, 特征向量2, ......], ......}
        self.label_feature_dict = dict()

    '''
    对模型进行数据训练
    1、获取特征值数量
    2、计算每个分类的概率p(c)
    3、将传入的训练特征向量数据和训练分类信息组织在一起
    
    train_data - 训练特征数据集
    train_labels - 训练特征数据对应的分类信息列表
    
    '''

    def train(self, train_data: list, train_labels: list):

        if len(train_data) != len(train_labels):
            raise ValueError("您输入的测试数组与类别数组的长度不一致！")

        # 获取测试数据特征值的数量
        self.feature_length = len(train_data[0])

        # 获取所有的类别的数量
        label_count = len(train_labels)

        # 获取一个不重复的类别集合
        distinct_labels = set(train_labels)

        # 遍历所有的类别，计算每个类别的概率p(c)
        for item in distinct_labels:
            self.label_probability_dict[item] = train_labels.count(item) / label_count
        '''
        zip函数：将两个数组交叉放置，比如：
        x1 = [a, b, c]
        x2 = [e, f, g]
        k = zip(x1, x2)
        k:[(a, e), (b,f), (c, g)]
        '''
        # 将传入的训练特征向量数据和训练分类信息组织在一起，放到一个dict里（self.label_feature_dict）
        # 格式为{'类别1':[特征向量1, 特征向量2, ......], '类别2':[特征向量1, 特征向量2, ......], ......}
        for feature, label in zip(train_data, train_labels):
            if label not in self.label_feature_dict:
                self.label_feature_dict[label] = []
            self.label_feature_dict[label].append(feature)

        return self

    '''
    对测试数据进行分类
    test_data - 测试数据
    all_labels - 所有的分类可选值
    '''

    def caculate(self, test_data, all_labels):
        if self.feature_length == -1:
            raise ValueError("您还没有进行数据训练，请先训练！")

        # 计算test_data分别为各个类别的可能性（概率）:{"类别1":概率1, "类别2":概率2.....}
        test_label_probability_dict = dict()
        for test_label in all_labels:
            p = 1

            # 获取当前类别的总体概率p(c)
            label_probability = self.label_probability_dict[test_label]

            # 当前类别的所有特征向量
            features_for_label = self.label_feature_dict[test_label]

            # 当前类别的记录数
            feature_record_count = len(features_for_label)

            # 转置一下，将相同维度的特征量放在一起，格式如：
            '''
            行列转置：
            原来的是 [t1, t2, t3, t4], [t1, t2, t3, t4]
            转置以后 [t1, t1], [t2, t2], [t3, t3], [t4, t4]
            
            features_for_label 转置后，行数为特征种类数量，每一行的列为所有记录的同一特征的数据
            '''
            features_for_label = numpy.array(features_for_label).T

            # 依次计算各特征的概率, len(test_data) 为特征数量，test_data是一维数组
            for index in range(0, len(test_data)):
                vector = list(features_for_label[index])
                # 获取当前类别中，特征值是当前值的概率，然后相乘
                p *= vector.count(test_data[index]) / feature_record_count
                # 如果test_data[0]:0, vector:[1, 0, 1, 0, 0, 1, 1], p为3/7
            # label_probability 相当于p(当前特征|c)
            test_label_probability_dict[test_label] = p * label_probability

        # 取出概率最大的一个类别
        most_probable_label = sorted(test_label_probability_dict,
                                     key=lambda x: test_label_probability_dict[x], reverse=True)[0]
        return most_probable_label
